#include "cminusf_builder.hpp"

// use these macros to get constant value
#define CONST_FP(num) \
    ConstantFP::get((float)num, module.get())
#define CONST_ZERO(type) \
    ConstantZero::get(type, module.get())


// You can define global variables here
// to store state
#define CONST_INT(num) \
    ConstantInt::get((int)num, module.get())
    
#define Int32Type \
 Type::get_int32_type(module.get())
 
#define FloatType \
 Type::get_float_type(module.get())

Value* Res;                     /* result */
Value* arg;                     
bool need_as_address = false;   /* val or addr */

/*
 * use CMinusfBuilder::Scope to construct scopes
 * scope.enter: enter a new scope
 * scope.exit: exit current scope
 * scope.push: add a new binding to current scope
 * scope.find: find and return the value bound to the name
 */

/*program->declaration-list */
void CminusfBuilder::visit(ASTProgram &node) {
 for (auto decl : node.declarations)         /* declaration-list */
     decl->accept(*this);                    
}

/* Num*/
void CminusfBuilder::visit(ASTNum &node) {
 if (node.type == TYPE_INT)                  
     Res = CONST_INT(node.i_val);   /* get int */
 else if (node.type == TYPE_FLOAT)
     Res = CONST_FP(node.f_val);    /* get float */
}

/*var-declaration -> type-specifier ID | type-specifier ID [INTEGER] */
void CminusfBuilder::visit(ASTVarDeclaration &node) {              
    Type* tmpType;                  /* get type */
    if (node.type == TYPE_INT)      
        tmpType = Int32Type;
    else if (node.type == TYPE_FLOAT)
        tmpType = FloatType;
        
    if (node.num != nullptr) {          /* is array*/
        auto* arrayType = ArrayType::get(tmpType, node.num->i_val);
        auto initializer = CONST_ZERO(tmpType);
        Value* arrayAlloca;             /* addr */
        if (scope.in_global())          /* global */
            arrayAlloca = GlobalVariable::create(node.id, module.get(), arrayType, false, initializer);
        else                            /* local */
            arrayAlloca = builder->create_alloca(arrayType);
        scope.push(node.id, arrayAlloca);
    }
    else {                              /* is not array */
        auto initializer = CONST_ZERO(tmpType);
        Value* varAlloca;               /* addr */
        if (scope.in_global())          /* global */
            varAlloca = GlobalVariable::create(node.id, module.get(), tmpType, false, initializer);
        else                            /* local */
            varAlloca = builder->create_alloca(tmpType);
        scope.push(node.id, varAlloca);
    }
}

/*fun-declaration -> type-specifier ID ( params ) compound-stmt */
void CminusfBuilder::visit(ASTFunDeclaration &node) {   
    Type* retType;               /* return type */
    if (node.type == TYPE_INT) { retType = Int32Type; }
    if (node.type == TYPE_FLOAT) { retType = FloatType; }
    if (node.type == TYPE_VOID) { retType = Type::get_void_type(module.get()); }

    std::vector<Type*> paramsType;  /* params type list */
    for (auto param : node.params) {
        if (param->isarray) {       /* is array */
            if (param->type == TYPE_INT)        /* int */
                paramsType.push_back(Type::get_int32_ptr_type(module.get()));
            else if(param->type == TYPE_FLOAT)  /* float */
                paramsType.push_back(Type::get_float_ptr_type(module.get()));
        }
        else {                  /* is not array */
            if (param->type == TYPE_INT)        /* int */
                paramsType.push_back(Int32Type);
            else if (param->type == TYPE_FLOAT) /* float */
                paramsType.push_back(FloatType);
        }
    }
    auto funType = FunctionType::get(retType, paramsType);              /* functiontype */
    auto function = Function::create(funType, node.id, module.get());   /* creat function */
    scope.push(node.id, function); 
    scope.enter();

    auto bb = BasicBlock::create(module.get(), node.id + "_entry", function);/* creat bb */
    builder->set_insert_point(bb);  /* insert bb into builder */

    std::vector<Value*> args;       /* args */
    for (auto arg = function->arg_begin();arg != function->arg_end();arg++) {
        args.push_back(*arg);
    }
    for (int i = 0;i < node.params.size();i++) {    /* params */
        auto param = node.params[i];
        arg = args[i];
        param->accept(*this);           /* accept */
    }

    node.compound_stmt->accept(*this);      /* compound-stmt */
    if (builder->get_insert_block()->get_terminator() == nullptr) {/* not terminator*/
        if (function->get_return_type()->is_void_type()) /*return void*/
            builder->create_void_ret();
        else if (function->get_return_type()->is_float_type())  /*return float*/
            builder->create_ret(CONST_FP(0.));
        else   /*return int*/
            builder->create_ret(CONST_INT(0));
    }
    scope.exit();                       
}

/*param→type-specifier ID ∣ type-specifier ID []*/
void CminusfBuilder::visit(ASTParam &node) {        
 Value* paramAlloca; /* addr space */
 if (node.isarray) { /* is array */
     if (node.type == TYPE_INT)        /* int */
         paramAlloca = builder->create_alloca(Type::get_int32_ptr_type(module.get()));
     else if (node.type == TYPE_FLOAT) /* float */
         paramAlloca = builder->create_alloca(Type::get_float_ptr_type(module.get()));
 }
 else {              /* is not array */
     if (node.type == TYPE_INT)        /* int*/
         paramAlloca = builder->create_alloca(Int32Type);
     else if (node.type == TYPE_FLOAT) /* int */
         paramAlloca = builder->create_alloca(FloatType);
 }
 builder->create_store(arg, paramAlloca);    /* load */
 scope.push(node.id, paramAlloca);          
}

/* CompoundStmt, compound-stmt -> {local-declarations statement-list} */
void CminusfBuilder::visit(ASTCompoundStmt &node) {
 scope.enter();     
 for (auto local_declaration : node.local_declarations)
     local_declaration->accept(*this);   /* local_declarations */
 for (auto statement : node.statement_list)
     statement->accept(*this);           /* statement */
 scope.exit(); 
}

/* ExpressionStmt, expression-stmt -> expression;
 *                                              | ; */
void CminusfBuilder::visit(ASTExpressionStmt &node) {
    if (node.expression != nullptr)     /* exist */
        node.expression->accept(*this);
}

/* SelectionStmt, selection-stmt -> if (expression) statement
 *                                        | if (expression) statement else statement */
void CminusfBuilder::visit(ASTSelectionStmt &node) { 
    auto func = builder->get_insert_block()->get_parent();/* current func */
    node.expression->accept(*this);        /* expression result */
    auto retType = Res->get_type();        /* type result */
    Value* ifExpression;                   /* if result */
    if(retType->is_integer_type())         /* int */
        ifExpression = builder->create_icmp_gt(Res, CONST_ZERO(Int32Type));                     /* >0:true */
    else                                   /* float */
        ifExpression = builder->create_fcmp_gt(Res, CONST_ZERO(FloatType));                     /* >0:true */

    auto ifBB = BasicBlock::create(module.get(), "ifBB", func);
    auto elseBB = BasicBlock::create(module.get(), "elseBB", func);
    auto afterIfBB = BasicBlock::create(module.get(), "afterIfBB", func);

    if(node.else_statement==nullptr)             /*no else*/
        builder->create_cond_br(ifExpression, ifBB, afterIfBB);
    else
        builder->create_cond_br(ifExpression, ifBB, elseBB);

    builder->set_insert_point(ifBB);            /*ifBB*/
    node.if_statement->accept(*this);
    if (builder->get_insert_block()->get_terminator() == nullptr)builder->create_br(afterIfBB);

    if(node.else_statement!=nullptr){          /*exist else*/
        builder->set_insert_point(elseBB);      /*elseBB*/
        node.else_statement->accept(*this);
        if (builder->get_insert_block()->get_terminator() == nullptr)builder->create_br(afterIfBB);        
    }
    else
        elseBB->erase_from_parent();          /*no else*/

    builder->set_insert_point(afterIfBB);
}
/* IterationStmt, iteration-stmt -> while (expression) statement */
void CminusfBuilder::visit(ASTIterationStmt &node) {    
    auto function = builder->get_insert_block()->get_parent();  /* current fun*/
    auto conditionBB = BasicBlock::create(module.get(), "condition", function); /* conditionBB */
    auto loopBB = BasicBlock::create(module.get(), "loop", function);           /* loopBB */
    auto retBB = BasicBlock::create(module.get(), "ret", function);             /* retBB */

    builder->create_br(conditionBB);        /* conditionBB */

    builder->set_insert_point(conditionBB);
    node.expression->accept(*this);         
    auto resType = Res->get_type();         /* return type */
    Value* TrueFalse;                       /* if expression */

    if (resType->is_integer_type()) {       /* int */
        auto intType = Int32Type;
        TrueFalse = builder->create_icmp_gt(Res, CONST_ZERO(intType));   /* >0:true */
    }
    else if (resType->is_float_type()) {    /* float */
        auto floatType = FloatType;
        TrueFalse = builder->create_fcmp_gt(Res, CONST_ZERO(floatType)); /* >0:true */
    }

    builder->create_cond_br(TrueFalse, loopBB, retBB);  /* cond */

    builder->set_insert_point(loopBB);      /* loopBB */
    node.statement->accept(*this);
    if (builder->get_insert_block()->get_terminator() == nullptr)builder->create_br(conditionBB);

    builder->set_insert_point(retBB);     /* retBB */
}
/* ReturnStmt, return-stmt -> return; | return expression; */
void CminusfBuilder::visit(ASTReturnStmt &node) {   
    auto function = builder->get_insert_block()->get_parent();  /* current fun */
    auto retType = function->get_return_type();     /* return type */
    if (retType->is_void_type()) {      /* void */
        builder->create_void_ret();     
        return;
    }

    node.expression->accept(*this);     /* expression */
    auto resType = Res->get_type();     /* expression result type*/

    if (retType->is_integer_type() && resType->is_float_type()) 
        Res = builder->create_fptosi(Res, Int32Type);
    if (retType->is_float_type() && resType->is_integer_type())
        Res = builder->create_sitofp(Res, FloatType);
    builder->create_ret(Res);           /* expression result*/
}

/* Var,var -> ID | ID [expression] */
void CminusfBuilder::visit(ASTVar &node) {     
    auto var = scope.find(node.id);             /* find var */
    bool addr_value = need_as_address;    /* addr_value */
    need_as_address = false;            /* reset need_as_address */
    Value* index = CONST_INT(0);        /* index=0 */
    if (node.expression != nullptr) {   /* expression */
        node.expression->accept(*this);
        auto res = Res;                 /* result */
        auto resType=res->get_type();  /* res type */
        if (res->get_type()==FloatType)            /* float */
            res = builder->create_fptosi(res, Int32Type);   /* change to int */
        index = res;                    /* index */

        auto function = builder->get_insert_block()->get_parent();  /* current fun */
        auto indexTest = builder->create_icmp_lt(index, CONST_ZERO(Int32Type)); /* index<0 */
        auto failBB = BasicBlock::create(module.get(), node.id + "_failbb", function);/* failbb */
        auto passBB = BasicBlock::create(module.get(), node.id + "_passbb", function);/* passbb */
        builder->create_cond_br(indexTest, failBB, passBB); /* br */

        builder->set_insert_point(failBB);  /* failBB */
        auto fail = scope.find("neg_idx_except");               /* find neg_idx_except */
        builder->create_call(static_cast<Function*>(fail), {}); /* neg_idx_except */
        builder->create_br(passBB);         /* goto passBB */

        builder->set_insert_point(passBB);  /* passBB */
        if (var->get_type()->get_pointer_element_type()->is_array_type())   /* is_array_pointer */
            var = builder->create_gep(var, { CONST_INT(0), index });    /* find addr */
        else {
            if (var->get_type()->get_pointer_element_type()->is_pointer_type()) /* is_pointer */
                var = builder->create_load(var);        /* addr */
            var = builder->create_gep(var, { index });  /* find var */
        }
        if (addr_value) {         /* return_addr */
            Res = var;                      
            need_as_address = false;        /* reset */
        }
        else 
            Res = builder->create_load(var);/* load */
        return;
    }
    /* no expression */
    if (addr_value) { /* return addr */
        Res = var;             
        need_as_address = false;/* reset */
    }
    else {                
        if (var->get_type()->get_pointer_element_type()->is_array_type())   /* is_array */
            Res = builder->create_gep(var, { CONST_INT(0), CONST_INT(0) }); /* find addr */
        else
            Res = builder->create_load(var);/* load */
    }
}

/* AssignExpression, 赋值语句, var = expression */
void CminusfBuilder::visit(ASTAssignExpression &node) {     
    need_as_address = true;         /* need val */
    node.var->accept(*this);        /* var */
    auto var = Res;                 /* addr */
    node.expression->accept(*this); /* expression */
    Value* res = Res;                 /* result */
    auto varType = var->get_type()->get_pointer_element_type();  /* var type */
    auto resType=res->get_type(); /* res type */
    if (varType == FloatType && resType == Int32Type)
        res = builder->create_sitofp(Res, FloatType);
    if (varType == Int32Type && resType == FloatType)
        res = builder->create_fptosi(Res, Int32Type);

    builder->create_store(res, var);/* load */
}

/* SimpleExpression, simple-expression -> additive-expression relop additive-expression | additive-expression */
void CminusfBuilder::visit(ASTSimpleExpression &node) {    
    node.additive_expression_l->accept(*this);  /* l expression */
    auto lres = Res;                            /* lres */
    if (node.additive_expression_r == nullptr) { return; }  //* no r expression */ 
    node.additive_expression_r->accept(*this);  /* r expression */
    auto rres = Res;                            /* rres */
    auto lresType=lres->get_type(); /* lres type */
    auto rresType=rres->get_type(); /* rres type */
    if (lresType == Int32Type && rresType == Int32Type) {     /* int int */
        switch (node.op) { 
        case OP_LE:
            Res = builder->create_icmp_le(lres, rres);break;
        case OP_LT:
            Res = builder->create_icmp_lt(lres, rres);break;
        case OP_GT:
            Res = builder->create_icmp_gt(lres, rres);break;
        case OP_GE:
            Res = builder->create_icmp_ge(lres, rres);break;
        case OP_EQ:
            Res = builder->create_icmp_eq(lres, rres);break;
        case OP_NEQ:
            Res = builder->create_icmp_ne(lres, rres);break;
        }
    }
    else {  /* float int or int float */
        if (lresType == Int32Type) /* int lres*/
            lres = builder->create_sitofp(lres, FloatType);
        if (rresType == Int32Type) /* int rres */
            rres = builder->create_sitofp(rres, FloatType);
        switch (node.op) {
        case OP_LE:
            Res = builder->create_fcmp_le(lres, rres);break;
        case OP_LT:
            Res = builder->create_fcmp_lt(lres, rres);break;
        case OP_GT:
            Res = builder->create_fcmp_gt(lres, rres);break;
        case OP_GE:
            Res = builder->create_fcmp_ge(lres, rres);break;
        case OP_EQ:
            Res = builder->create_fcmp_eq(lres, rres);break;
        case OP_NEQ:
            Res = builder->create_fcmp_ne(lres, rres);break;
        }
    }
    Res = builder->create_zext(Res, Int32Type); 
}
/* AdditiveExpression,additive-expression -> additive-expression addop term | term */
void CminusfBuilder::visit(ASTAdditiveExpression &node) {   
    if (node.additive_expression == nullptr) {  /* no addop */
        node.term->accept(*this);return;
    }
    node.additive_expression->accept(*this);    /* expression */
    auto lres = Res;                            /* lres */
    node.term->accept(*this);                   /* term */
    auto rres = Res;                            /* rres */
    auto lresType=lres->get_type(); /* lres type */
    auto rresType=rres->get_type(); /* rres type */
    if (lresType == Int32Type && rresType == Int32Type) {     /* int int */
        switch (node.op) { 
        case OP_PLUS:
            Res = builder->create_iadd(lres, rres);break;
        case OP_MINUS:
            Res = builder->create_isub(lres, rres);break;
        }
    }
    else {  /* float int or int float */
        if (lresType == Int32Type) /* int lres */
            lres = builder->create_sitofp(lres, FloatType);
        if (rresType == Int32Type) /* int rres */
            rres = builder->create_sitofp(rres, FloatType);
        switch (node.op) { 
        case OP_PLUS:
            Res = builder->create_fadd(lres, rres);break;
        case OP_MINUS:
            Res = builder->create_fsub(lres, rres);break;
        }
    }
}


/* Term, Term -> term mulop factor | factor */
void CminusfBuilder::visit(ASTTerm &node) {     
    if (node.term == nullptr) {             /* no term */
        node.factor->accept(*this);return;
    }
    node.term->accept(*this);   /* left */
    auto lres = Res;            /* lres */
    node.factor->accept(*this); /* right */
    auto rres = Res;            /* rres */
    auto lresType=lres->get_type(); /* lres type */
    auto rresType=rres->get_type(); /* rres type */
    if (lresType == Int32Type && rresType == Int32Type) { /* int int */
        switch (node.op) { 
        case OP_MUL:
            Res = builder->create_imul(lres, rres);break;
        case OP_DIV:
            Res = builder->create_isdiv(lres, rres);break;
        }
    }
    else {  /* float int or int float */
        if (lresType == Int32Type) /* int lres */
            lres = builder->create_sitofp(lres, FloatType);
        if (rresType == Int32Type) /* int rres */
            rres = builder->create_sitofp(rres, FloatType);
        switch (node.op) { 
        case OP_MUL:
            Res = builder->create_fmul(lres, rres);break;
        case OP_DIV:
            Res = builder->create_fdiv(lres, rres);break;
        }
    }
}


/* Call, 函数调用, call -> ID (args) */
void CminusfBuilder::visit(ASTCall &node) {    
    auto function = static_cast<Function*>(scope.find(node.id));    /* find fun */
    auto paramType = function->get_function_type()->param_begin();  /* get param trpe */
    std::vector<Value*> args;       /* args */
    for (auto arg : node.args) {    /* arg */
        arg->accept(*this);         /* get res */
        if (Res->get_type()->is_pointer_type()) {   /* is pointer */
            args.push_back(Res);    
        }
        else {  /* is not pointer */
            auto resType=Res->get_type(); /* res type */
            if (resType == FloatType && resType == Int32Type)
                Res = builder->create_sitofp(Res, FloatType);
            if (resType == Int32Type && resType == FloatType)
                Res = builder->create_fptosi(Res, Int32Type);
            args.push_back(Res);
        }
        paramType++;                /* next param */
    }
    Res = builder->create_call(static_cast<Function*>(function), args); /* create call */
}
